from copy import deepcopy

from lagent.schema import ActionReturn, ActionStatusCode, FunctionCall
from .hook import Hook


class ActionPreprocessor(Hook):
    """
    动作预处理器钩子
    
    用于预处理动作消息和后处理动作返回消息的钩子类
    继承自基础Hook类
    """

    def before_action(self, executor, message, session_id):
        """
        动作执行前的预处理
        
        验证消息格式并提取动作名称和参数
        
        参数:
            executor: 执行器实例
            message: 待处理的消息
            session_id: 会话ID
            
        返回:
            处理后的消息
        """
        # 验证消息格式是否符合以下三种情况之一:
        # 1. FunctionCall实例
        # 2. 包含name和parameters的字典
        # 3. 包含action字典且action中有name和parameters的字典
        assert isinstance(message.formatted, FunctionCall) or (
            isinstance(message.formatted, dict) and 'name' in message.content
            and 'parameters' in message.formatted) or (
                'action' in message.formatted
                and 'parameters' in message.formatted['action']
                and 'name' in message.formatted['action'])

        # 提取动作名称和参数
        if isinstance(message.formatted, dict):
            name = message.formatted.get('name',
                                       message.formatted['action']['name'])
            parameters = message.formatted.get(
                'parameters', message.formatted['action']['parameters'])
        else:
            name = message.formatted.name
            parameters = message.formatted.parameters
            
        # 更新消息内容
        message.content = dict(name=name, parameters=parameters)
        return message

    def after_action(self, executor, message, session_id):
        """
        动作执行后的后处理
        
        格式化动作执行结果
        
        参数:
            executor: 执行器实例
            message: 待处理的消息
            session_id: 会话ID
            
        返回:
            处理后的消息
        """
        action_return = message.content
        
        # 处理ActionReturn类型的返回结果
        if isinstance(action_return, ActionReturn):
            if action_return.state == ActionStatusCode.SUCCESS:
                response = action_return.format_result()  # 成功则格式化结果
            else:
                response = action_return.errmsg  # 失败则返回错误信息
        else:
            response = action_return  # 其他情况直接返回
            
        message.content = response
        return message


class InternLMActionProcessor(ActionPreprocessor):
    """
    InternLM动作处理器
    
    专门用于处理InternLM模型动作的预处理器
    继承自ActionPreprocessor
    """

    def __init__(self, code_parameter: str = 'command'):
        """
        初始化处理器
        
        参数:
            code_parameter (str): 代码参数名，默认为'command'
        """
        self.code_parameter = code_parameter

    def before_action(self, executor, message, session_id):
        """
        动作执行前的预处理
        
        处理InternLM特定的消息格式
        
        参数:
            executor: 执行器实例
            message: 待处理的消息
            session_id: 会话ID
            
        返回:
            处理后的消息
        """
        # 创建消息的深拷贝以避免修改原始消息
        message = deepcopy(message)
        
        # 验证消息格式
        assert isinstance(message.formatted, dict) and set(
            message.formatted).issuperset(
                {'tool_type', 'thought', 'action', 'status'})
                
        # 处理字符串类型的动作
        if isinstance(message.formatted['action'], str):
            # 封装代码解释器参数
            action_name = next(iter(executor.actions))
            parameters = {self.code_parameter: message.formatted['action']}
            
            # 为异步解释器添加会话ID
            if action_name in ['AsyncIPythonInterpreter']:
                parameters['session_id'] = session_id
                
            # 更新动作格式
            message.formatted['action'] = dict(
                name=action_name, parameters=parameters)
                
        # 调用父类的预处理方法
        return super().before_action(executor, message, session_id)
